"""
HuaTuoGPT 模型下载脚本
使用镜像加速下载
"""
import os
from transformers import AutoTokenizer, AutoModelForCausalLM

# 使用国内镜像加速
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'

model_name = "FreedomIntelligence/HuatuoGPT-II-7B"

print("=" * 80)
print("🚀 开始下载 HuaTuoGPT 模型")
print("=" * 80)
print(f"模型: {model_name}")
print(f"大小: 约 14GB")
print(f"缓存位置: {os.path.expanduser('~/.cache/huggingface/')}")
print("")
print("⏳ 下载中，请耐心等待（可能需要 10-30 分钟）...")
print("")

try:
    # 下载 Tokenizer
    print("📥 步骤 1/2: 下载 Tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(
        model_name,
        trust_remote_code=True
    )
    print("✅ Tokenizer 下载完成")
    
    # 下载模型
    print("")
    print("📥 步骤 2/2: 下载模型文件（约 14GB）...")
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        trust_remote_code=True,
        torch_dtype="auto",
        low_cpu_mem_usage=True
    )
    print("✅ 模型下载完成")
    
    print("")
    print("=" * 80)
    print("🎉 HuaTuoGPT 模型下载成功！")
    print("=" * 80)
    print(f"缓存位置: {os.path.expanduser('~/.cache/huggingface/')}")
    print("")
    print("下一步：运行 python verify_huatuogpt.py 验证安装")
    print("")
    
except Exception as e:
    print("")
    print("=" * 80)
    print(f"❌ 下载失败: {e}")
    print("=" * 80)
    print("")
    print("可能的原因：")
    print("1. 网络连接问题")
    print("2. 磁盘空间不足（需要至少 20GB）")
    print("3. 防火墙阻止")
    print("")
    print("建议：")
    print("1. 检查网络连接")
    print("2. 检查磁盘空间")
    print("3. 参考 HuaTuoGPT手动部署方案.md 中的故障排除部分")
    print("")

